########################################################################
#
# 双目相机标定程序
# 使用img文件夹下的标定图片进行相机标定
#
########################################################################

import os
import cv2
import numpy as np
import json
import glob
from pathlib import Path

# 设置numpy显示格式，取消科学计数法
np.set_printoptions(suppress=True, precision=6)


class StereoCameraCalibrator:
    """双目相机标定器"""
    
    def __init__(self, img_folder="img", chessboard_size=(9, 6), square_size=25.0):
        """
        初始化标定器
        
        Args:
            img_folder: 标定图片文件夹路径
            chessboard_size: 棋盘格内角点数量 (width, height)
            square_size: 棋盘格方格大小(mm)
        """
        self.img_folder = img_folder
        self.chessboard_size = chessboard_size
        self.square_size = square_size
        
        # 标定结果
        self.left_camera_matrix = None
        self.left_dist_coeffs = None
        self.right_camera_matrix = None
        self.right_dist_coeffs = None
        self.R = None  # 旋转矩阵
        self.T = None  # 平移向量
        self.E = None  # 本质矩阵
        self.F = None  # 基础矩阵
        
        # 标定质量指标
        self.left_reprojection_error = None
        self.right_reprojection_error = None
        self.stereo_reprojection_error = None
        
        # 准备棋盘格世界坐标
        self.prepare_object_points()
    
    def prepare_object_points(self):
        """准备棋盘格的世界坐标"""
        # 创建棋盘格的世界坐标
        objp = np.zeros((self.chessboard_size[0] * self.chessboard_size[1], 3), np.float32)
        objp[:, :2] = np.mgrid[0:self.chessboard_size[0], 0:self.chessboard_size[1]].T.reshape(-1, 2)
        objp *= self.square_size  # 转换为实际尺寸(mm)
        self.objp = objp
    
    def find_calibration_images(self):
        """查找标定图片"""
        # 查找左右相机图片
        left_pattern = os.path.join(self.img_folder, "left*.jpg")
        right_pattern = os.path.join(self.img_folder, "right*.jpg")
        
        left_images = sorted(glob.glob(left_pattern))
        right_images = sorted(glob.glob(right_pattern))
        
        print(f"找到 {len(left_images)} 张左相机图片")
        print(f"找到 {len(right_images)} 张右相机图片")
        
        # 确保左右图片数量一致
        min_count = min(len(left_images), len(right_images))
        left_images = left_images[:min_count]
        right_images = right_images[:min_count]
        
        return left_images, right_images
    
    def detect_chessboard_corners(self, image_path):
        """检测棋盘格角点"""
        image = cv2.imread(image_path)
        if image is None:
            print(f"无法读取图片: {image_path}")
            return None, None
        
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        
        # 查找棋盘格角点
        ret, corners = cv2.findChessboardCorners(gray, self.chessboard_size, None)
        
        if ret:
            # 亚像素角点检测
            criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
            corners = cv2.cornerSubPix(gray, corners, (11, 11), (-1, -1), criteria)
            return corners, image
        else:
            print(f"未能在图片中检测到棋盘格: {image_path}")
            return None, image
    
    def calibrate_single_camera(self, image_paths, camera_name):
        """标定单个相机"""
        print(f"\n开始标定 {camera_name} 相机...")
        
        objpoints = []  # 世界坐标点
        imgpoints = []  # 图像坐标点
        valid_images = []
        
        for i, img_path in enumerate(image_paths):
            print(f"处理图片 {i+1}/{len(image_paths)}: {os.path.basename(img_path)}")
            
            corners, image = self.detect_chessboard_corners(img_path)
            if corners is not None:
                objpoints.append(self.objp)
                imgpoints.append(corners)
                valid_images.append(img_path)
                
                # 绘制角点（可选）
                cv2.drawChessboardCorners(image, self.chessboard_size, corners, True)
                cv2.imshow(f'{camera_name} 角点检测', image)
                cv2.waitKey(500)
        
        cv2.destroyAllWindows()
        
        if len(objpoints) < 3:
            print(f"有效图片数量不足，至少需要3张图片，当前只有 {len(objpoints)} 张")
            return None, None, None
        
        print(f"使用 {len(objpoints)} 张有效图片进行标定")
        
        # 相机标定
        ret, camera_matrix, dist_coeffs, rvecs, tvecs = cv2.calibrateCamera(
            objpoints, imgpoints, image.shape[:2][::-1], None, None
        )
        
        if ret:
            print(f"✓ {camera_name} 相机标定成功")
            
            # 计算重投影误差
            mean_error = 0
            for i in range(len(objpoints)):
                imgpoints2, _ = cv2.projectPoints(objpoints[i], rvecs[i], tvecs[i], 
                                                camera_matrix, dist_coeffs)
                error = cv2.norm(imgpoints[i], imgpoints2, cv2.NORM_L2) / len(imgpoints2)
                mean_error += error
            
            reprojection_error = mean_error / len(objpoints)
            print(f"重投影误差: {reprojection_error:.4f} 像素")
            
            return camera_matrix, dist_coeffs, reprojection_error
        else:
            print(f"✗ {camera_name} 相机标定失败")
            return None, None, None
    
    def calibrate_stereo_camera(self, left_images, right_images):
        """标定双目相机"""
        print("\n开始双目相机标定...")
        
        # 标定左右相机
        left_camera_matrix, left_dist_coeffs, left_error = self.calibrate_single_camera(
            left_images, "左相机"
        )
        right_camera_matrix, right_dist_coeffs, right_error = self.calibrate_single_camera(
            right_images, "右相机"
        )
        
        if left_camera_matrix is None or right_camera_matrix is None:
            print("单相机标定失败，无法进行双目标定")
            return False
        
        # 准备双目标定数据
        objpoints = []
        left_imgpoints = []
        right_imgpoints = []
        
        for left_path, right_path in zip(left_images, right_images):
            left_corners, _ = self.detect_chessboard_corners(left_path)
            right_corners, _ = self.detect_chessboard_corners(right_path)
            
            if left_corners is not None and right_corners is not None:
                objpoints.append(self.objp)
                left_imgpoints.append(left_corners)
                right_imgpoints.append(right_corners)
        
        if len(objpoints) < 3:
            print("有效图片对数量不足，无法进行双目标定")
            return False
        
        print(f"使用 {len(objpoints)} 对有效图片进行双目标定")
        
        # 双目相机标定
        ret, left_camera_matrix, left_dist_coeffs, right_camera_matrix, right_dist_coeffs, \
        R, T, E, F = cv2.stereoCalibrate(
            objpoints, left_imgpoints, right_imgpoints,
            left_camera_matrix, left_dist_coeffs,
            right_camera_matrix, right_dist_coeffs,
            cv2.imread(left_images[0]).shape[:2][::-1],
            flags=cv2.CALIB_FIX_INTRINSIC
        )
        
        if ret:
            print("✓ 双目相机标定成功")
            
            # 计算立体标定重投影误差
            mean_error = 0
            for i in range(len(objpoints)):
                left_imgpoints2, _ = cv2.projectPoints(objpoints[i], 
                                                     cv2.Rodrigues(R)[0], T,
                                                     left_camera_matrix, left_dist_coeffs)
                right_imgpoints2, _ = cv2.projectPoints(objpoints[i],
                                                      cv2.Rodrigues(R)[0], T,
                                                      right_camera_matrix, right_dist_coeffs)
                
                left_error = cv2.norm(left_imgpoints[i], left_imgpoints2, cv2.NORM_L2) / len(left_imgpoints2)
                right_error = cv2.norm(right_imgpoints[i], right_imgpoints2, cv2.NORM_L2) / len(right_imgpoints2)
                mean_error += (left_error + right_error) / 2
            
            stereo_error = mean_error / len(objpoints)
            print(f"立体标定重投影误差: {stereo_error:.4f} 像素")
            
            # 保存标定结果
            self.left_camera_matrix = left_camera_matrix
            self.left_dist_coeffs = left_dist_coeffs
            self.right_camera_matrix = right_camera_matrix
            self.right_dist_coeffs = right_dist_coeffs
            self.R = R
            self.T = T
            self.E = E
            self.F = F
            self.left_reprojection_error = left_error
            self.right_reprojection_error = right_error
            self.stereo_reprojection_error = stereo_error
            
            return True
        else:
            print("✗ 双目相机标定失败")
            return False
    
    def save_calibration_results(self, output_file="camera_calibration.json"):
        """保存标定结果到JSON文件"""
        if self.left_camera_matrix is None:
            print("没有标定结果可保存")
            return False
        
        # 设置numpy显示格式，取消科学计数法
        np.set_printoptions(suppress=True, precision=6)
        
        calibration_data = {
            "calibration_info": {
                "chessboard_size": self.chessboard_size,
                "square_size_mm": self.square_size,
                "calibration_date": str(Path().cwd()),
                "img_folder": self.img_folder
            },
            "left_camera": {
                "camera_matrix": self.left_camera_matrix.tolist(),
                "dist_coeffs": self.left_dist_coeffs.tolist(),
                "reprojection_error": float(self.left_reprojection_error)
            },
            "right_camera": {
                "camera_matrix": self.right_camera_matrix.tolist(),
                "dist_coeffs": self.right_dist_coeffs.tolist(),
                "reprojection_error": float(self.right_reprojection_error)
            },
            "stereo_camera": {
                "rotation_matrix": self.R.tolist(),
                "translation_vector": self.T.tolist(),
                "essential_matrix": self.E.tolist(),
                "fundamental_matrix": self.F.tolist(),
                "stereo_reprojection_error": float(self.stereo_reprojection_error)
            }
        }
        
        try:
            with open(output_file, 'w', encoding='utf-8') as f:
                json.dump(calibration_data, f, indent=2, ensure_ascii=False)
            print(f"✓ 标定结果已保存到: {output_file}")
            return True
        except Exception as e:
            print(f"✗ 保存标定结果失败: {e}")
            return False
    
    def print_calibration_summary(self):
        """打印标定结果摘要"""
        if self.left_camera_matrix is None:
            print("没有标定结果可显示")
            return
        
        # 设置numpy显示格式，取消科学计数法
        np.set_printoptions(suppress=True, precision=6)
        
        print("\n" + "="*50)
        print("标定结果摘要")
        print("="*50)
        
        print(f"棋盘格尺寸: {self.chessboard_size[0]} x {self.chessboard_size[1]}")
        print(f"方格大小: {self.square_size} mm")
        print(f"图片文件夹: {self.img_folder}")
        
        print(f"\n左相机重投影误差: {self.left_reprojection_error:.6f} 像素")
        print(f"右相机重投影误差: {self.right_reprojection_error:.6f} 像素")
        print(f"立体标定重投影误差: {self.stereo_reprojection_error:.6f} 像素")
        
        print(f"\n左相机内参矩阵:")
        print(self.left_camera_matrix)
        print(f"\n右相机内参矩阵:")
        print(self.right_camera_matrix)
        
        print(f"\n旋转矩阵:")
        print(self.R)
        print(f"\n平移向量:")
        print(self.T)
        
        print("="*50)
    
    def run_calibration(self):
        """运行完整的标定流程"""
        print("双目相机标定程序")
        print("="*50)
        
        # 检查图片文件夹
        if not os.path.exists(self.img_folder):
            print(f"图片文件夹不存在: {self.img_folder}")
            return False
        
        # 查找标定图片
        left_images, right_images = self.find_calibration_images()
        if len(left_images) == 0 or len(right_images) == 0:
            print("未找到标定图片")
            return False
        
        # 进行双目标定
        success = self.calibrate_stereo_camera(left_images, right_images)
        
        if success:
            # 打印标定结果
            self.print_calibration_summary()
            
            # 保存标定结果到脚本所在目录
            script_dir = os.path.dirname(os.path.abspath(__file__))
            output_file = os.path.join(script_dir, "camera_calibration.json")
            self.save_calibration_results(output_file)
            
            return True
        else:
            print("标定失败")
            return False


def main():
    """主函数"""
    # 获取当前脚本所在目录
    script_dir = os.path.dirname(os.path.abspath(__file__))
    img_folder = os.path.join(script_dir, "img")
    output_file = os.path.join(script_dir, "camera_calibration.json")
    
    print(f"脚本目录: {script_dir}")
    print(f"图片文件夹路径: {img_folder}")
    print(f"标定结果保存路径: {output_file}")
    
    # 创建标定器
    calibrator = StereoCameraCalibrator(
        img_folder=img_folder,
        chessboard_size=(11, 15),  # 棋盘格内角点数量
        square_size=30  # 方格大小(mm)
    )
    
    # 运行标定
    success = calibrator.run_calibration()
    
    if success:
        print(f"\n标定完成！结果已保存到: {output_file}")
    else:
        print("\n标定失败！")


if __name__ == "__main__":
    main()
